import cv2
import numpy as np
from keras.layers import activations, Conv2D, MaxPool2D, Dense, Flatten, Input, Permute, Reshape
from keras.layers.advanced_activations import PReLU
from keras.models import Model, Sequential
from utils import utils


# 创建p-net,注意pnet的训练数据仅为(12, 12 ,3),而测试数据shape可以任意，因为是全卷积层，所以不用随机截取图像获得候选框.
def build_pnet(path):
    # 输入 [n m 3]
    inputs = Input(shape=[None, None, 3])

    # -> c1[n-2 m-2 10]
    x = Conv2D(10, (3, 3), strides=(1, 1), padding='valid', name='conv1')(inputs)
    # shared_axes表示pRelu的参数两个轴共用。
    x = PReLU(shared_axes=[1, 2], name='PReLU1')(x)

    # -> p1[(n-2)/2 (m-2)/2 10] -> c2[(n-2)/2-2  ...
    x = MaxPool2D(pool_size=(2, 2), name='p1')(x)
    x = Conv2D(16, (3, 3), strides=1, padding='valid', name='conv2')(x)
    x = PReLU(shared_axes=[1, 2], name='PReLU2')(x)

    x = Conv2D(32, (3, 3), strides=1, padding='valid', name='conv3')(x)
    x = PReLU(shared_axes=[1, 2], name='PReLU3')(x)

    # 分类层
    classifier = Conv2D(2, (1, 1), activation='softmax', name='conv4-1')(x)

    # 框回归层
    bbox_regress = Conv2D(4, (1, 1), name='conv4-2')(x)

    model = Model([inputs], [classifier, bbox_regress])
    model.load_weights(path, by_name=True)
    return model


# r-net和o-net的输入固定分别为24 24 3和48 48 3
def build_rnet(path):
    # 输入[24 24 3]
    inputs = Input(shape=[24, 24, 3])

    # ->c1[22 22 28] -> p1[11 11 28]
    x = Conv2D(28, (3, 3), strides=1, padding='valid', name='conv1')(inputs)
    x = PReLU(shared_axes=[1, 2], name='prelu1')(x)
    # 注意池化的strides=2，且要补0。优先补左上。
    x = MaxPool2D(pool_size=(3, 3), strides=2, padding='same', name='p1')(x)

    # -> c2[9 9 48] -> p2[4 4 48]
    x = Conv2D(48, (3, 3), strides=1, padding='valid', name='conv2')(x)
    x = PReLU(shared_axes=[1, 2], name='prelu2')(x)
    x = MaxPool2D(pool_size=(3, 3), strides=2, name='p2')(x)

    # -> c3[3 3 64]
    x = Conv2D(64, (2, 2), strides=1, padding='valid', name='conv3')(x)
    x = PReLU(shared_axes=[1, 2], name='prelu3')(x)

    # 3 3 64 -> 64 3 3 ->576 切换索引顺序并展平，注意索引从1开始
    x = Permute(dims=(3, 2, 1))(x)
    x = Flatten()(x)

    # 576->128
    x = Dense(128, name='conv4')(x)
    x = PReLU(name='prelu4')(x)

    # 分类得分与框回归偏移量
    # 128->2
    classifier = Dense(2, activation='softmax', name='conv5-1')(x)
    bboxRegress = Dense(4, name='conv5-2')(x)

    model = Model([inputs], [classifier, bboxRegress])
    model.load_weights(path, by_name=True)
    return model


def build_onet(path):
    # 输入 [48 48 3]
    inputs = Input(shape=[48, 48, 3])

    # -> c1[46 46 32] -> p1[23 23 32]
    x = Conv2D(32, (3, 3), strides=1, padding='valid', name='conv1')(inputs)
    x = PReLU(shared_axes=[1, 2], name='prelu1')(x)
    x = MaxPool2D(pool_size=3, strides=2, padding='same', name='p1')(x)

    # -> c2[21 21 64] -> p2[10 10 64]
    x = Conv2D(64, (3, 3), strides=1, padding='valid', name='conv2')(x)
    x = PReLU(shared_axes=[1, 2], name='prelu2')(x)
    x = MaxPool2D(pool_size=3, strides=2, padding='valid', name='p2')(x)

    # -> c3[8 8 64] -> p3[4 4 64]
    x = Conv2D(64, (3, 3), strides=1, padding='valid', name='conv3')(x)
    x = PReLU(shared_axes=[1, 2], name='prelu3')(x)
    x = MaxPool2D(pool_size=2, name='p3')(x)

    # -> c4[3 3 128]
    x = Conv2D(128, (2, 2), strides=1, padding='valid', name='conv4')(x)
    x = PReLU(shared_axes=[1, 2], name='prelu4')(x)

    # ->permute[128 3 3] -> fla[1152]
    x = Permute((3, 2, 1))(x)
    x = Flatten()(x)

    # 1152->256
    x = Dense(256, name='conv5')(x)
    x = PReLU(name='prelu5')(x)

    # 结果层 256-> 2/4/10
    classifier = Dense(2, activation='softmax', name='conv6-1')(x)
    bboxRegress = Dense(4, name='conv6-2')(x)
    landmarkRegress = Dense(10, name='conv6-3')(x)

    model = Model([inputs], [classifier, bboxRegress, landmarkRegress])
    model.load_weights(path, by_name=True)
    return model


class mtcnn():
    def __init__(self):
        # 已经训练好的模型。
        self.oNet = build_onet('model_data/onet.h5')
        self.pNet = build_pnet('model_data/pnet.h5')
        self.rNet = build_rnet('model_data/rnet.h5')

    # 调用人脸检测框和关键点函数 threshould是每个网络人脸分类得分的阈值。取(0.5,0.7,0.9)
    def detect_face(self, img, threshold):
        # 对图像归一化，加快收敛速度，训练和测试都需要
        normalImg = (img.copy() - 127.5) / 127.5
        originH, originW, tmp = img.shape

        # 对p-net的结果进行解码处理得到对应原图的先验框
        prBoxes = []
        # 获得缩放的大小
        scales = utils.get_scales(originH, originW)
        for scale in scales:
            h = int(originH * scale)
            w = int(originW * scale)
            # 缩放图像并输入进pnet
            scaleImg = cv2.resize(normalImg, (w, h))
            # 输入多增一个维度？注意后面的out前面有list和多一维，两个维度
            inImg = np.expand_dims(scaleImg, 0)
            # 这里pnetOut为两个list,pnetOut[0][0]代表了分类结果, pnetOut[1][0]代表框偏移量预测结果。
            pnetOut = self.pNet.predict(inImg)
            # 对网络输出处理，得到num_box*5的array。
            prBox = utils.process_pnet_12out(pnetOut, scale, threshold[0], originH, originW)
            # 注意extend后是list,若要运算则需转化为np.array。
            prBoxes.extend(prBox)

        if len(prBoxes)==0:
            return prBoxes
        # NMS
        prBoxes = np.array(utils.NMS(np.array(prBoxes), 0.7))

        # rnet
        scaleImgs = []
        for box in prBoxes:
            # 剪切并resize
            cropImg = normalImg[int(box[1]):int(box[3]), int(box[0]):int(box[2])]
            scaleImgs.append(cv2.resize(cropImg, (24, 24)))
        claOut, offset = self.rNet.predict(np.array(scaleImgs))
        rnetBoxes = utils.process_rnet_24out(claOut, offset, prBoxes, threshold[1], originH, originW)
        if len(rnetBoxes)==0:
            return rnetBoxes

        # onet
        scaleImgs = []
        for box in rnetBoxes:
            # 剪切并resize
            cropImg = normalImg[int(box[1]):int(box[3]), int(box[0]):int(box[2])]
            scaleImgs.append(cv2.resize(cropImg, (48, 48)))
        claOut, offset, landmarkRegress = self.oNet.predict(np.array(scaleImgs))
        onetBoxes = utils.process_onet_48out(claOut, offset, landmarkRegress, rnetBoxes, threshold[2], originH, originW)
        return onetBoxes
